#include "dexterous_hand.h"
#include <crow/app.h>
#include <crow/websocket.h>
#include <nlohmann/json.hpp>
#include <iostream>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>
#include <queue>
#include <condition_variable>
#include <map>
#include <memory>
#include <array>
#include <cmath>
#include <cstdlib>

#ifdef _WIN32
#include <windows.h>
#include <io.h>
#include <fcntl.h>
#endif

using json = nlohmann::json;

// === 日志工具函数 ===
class Logger {
public:
    static void log(const std::string& level, const std::string& message) {
        auto now = std::chrono::system_clock::now();
        auto time_t = std::chrono::system_clock::to_time_t(now);
        auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
            now.time_since_epoch()) % 1000;
        
        std::tm tm_now;
#if defined(_WIN32)
        localtime_s(&tm_now, &time_t);
#else
        localtime_r(&time_t, &tm_now);
#endif
        
        char time_buffer[100];
        std::strftime(time_buffer, sizeof(time_buffer), "%Y-%m-%d %H:%M:%S", &tm_now);
        
        std::cout << "[" << time_buffer << "." 
                  << std::setfill('0') << std::setw(3) << ms.count() 
                  << "] [" << level << "] " << message << std::endl;
    }
    
    static void info(const std::string& message) { log("INFO", message); }
    static void warn(const std::string& message) { log("WARN", message); }
    static void error(const std::string& message) { log("ERROR", message); }
    static void debug(const std::string& message) { log("DEBUG", message); }
};

// === 线程安全的数据缓冲区模板 ===
template<typename T>
class ThreadSafeBuffer {
private:
    T data_;
    mutable std::mutex mtx_;
    std::condition_variable cv_;
    bool has_data_{false};

public:
    void set(const T& value) {
        std::lock_guard<std::mutex> lock(mtx_);
        data_ = value;
        has_data_ = true;
        cv_.notify_all();
    }

    T get() const {
        std::lock_guard<std::mutex> lock(mtx_);
        return data_;
    }

    bool wait_for_data(std::chrono::milliseconds timeout = std::chrono::milliseconds(100)) {
        std::unique_lock<std::mutex> lock(mtx_);
        return cv_.wait_for(lock, timeout, [this] { return has_data_; });
    }
};

// === 关节参数数组的包装器 ===
struct JointCommandsArray {
    JointParamsDownlink joints[21];
    
    JointCommandsArray() {
        for (int i = 0; i < 21; ++i) {
            joints[i] = {};
        }
    }
    
    JointCommandsArray(const JointCommandsArray& other) {
        for (int i = 0; i < 21; ++i) {
            joints[i] = other.joints[i];
        }
    }
    
    JointCommandsArray& operator=(const JointCommandsArray& other) {
        if (this != &other) {
            for (int i = 0; i < 21; ++i) {
                joints[i] = other.joints[i];
            }
        }
        return *this;
    }
};

// === 健康监控类 ===
class HealthMonitor {
private:
    std::atomic<int> serial_error_count_;
    std::atomic<uint64_t> last_successful_read_ms_;
    std::atomic<uint64_t> total_messages_sent_;
    std::atomic<uint64_t> total_messages_received_;

public:
    HealthMonitor() : serial_error_count_(0), last_successful_read_ms_(0), 
                      total_messages_sent_(0), total_messages_received_(0) {}

    void reportSerialSuccess() {
        serial_error_count_ = 0;
        last_successful_read_ms_ = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now().time_since_epoch()).count();
        total_messages_received_++;
    }

    void reportSerialError() {
        serial_error_count_++;
    }

    void reportMessageSent() {
        total_messages_sent_++;
    }

    bool isHealthy() const {
        auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now().time_since_epoch()).count();
        auto elapsed = now_ms - last_successful_read_ms_.load();
        
        return serial_error_count_ < 10 && elapsed < 1000;
    }

    void printStats() const {
        std::ostringstream oss;
        oss << "Serial errors: " << serial_error_count_ 
            << ", Messages sent: " << total_messages_sent_ 
            << ", Messages received: " << total_messages_received_ 
            << ", Status: " << (isHealthy() ? "HEALTHY" : "UNHEALTHY");
        Logger::info(oss.str());
    }
};

// === 安全的WebSocket客户端管理器 ===
class SafeClientManager {
private:
    std::map<uint64_t, crow::websocket::connection*> clients_;
    mutable std::mutex clients_mtx_;
    std::atomic<uint64_t> next_id_;

public:
    SafeClientManager() : next_id_(0) {}

    uint64_t addClient(crow::websocket::connection* conn) {
        std::lock_guard<std::mutex> lock(clients_mtx_);
        uint64_t id = ++next_id_;
        clients_[id] = conn;
        
        std::ostringstream oss;
        oss << "WebSocket client connected - ID: " << id << ", Total clients: " << clients_.size();
        Logger::info(oss.str());
        
        return id;
    }

    void removeClient(uint64_t id) {
        std::lock_guard<std::mutex> lock(clients_mtx_);
        if (clients_.erase(id)) {
            std::ostringstream oss;
            oss << "WebSocket client disconnected - ID: " << id << ", Total clients: " << clients_.size();
            Logger::info(oss.str());
        }
    }

    void broadcast(const std::string& message) {
        std::vector<uint64_t> failed_clients;
        size_t success_count = 0;
        
        {
            std::lock_guard<std::mutex> lock(clients_mtx_);
            for (auto it = clients_.begin(); it != clients_.end(); ++it) {
                uint64_t id = it->first;
                crow::websocket::connection* conn = it->second;
                try {
                    conn->send_text(message);
                    success_count++;
                } catch (const std::exception& e) {
                    std::ostringstream oss;
                    oss << "Failed to send message to client ID " << id << ": " << e.what();
                    Logger::error(oss.str());
                    failed_clients.push_back(id);
                } catch (...) {
                    std::ostringstream oss;
                    oss << "Failed to send message to client ID " << id << ": Unknown error";
                    Logger::error(oss.str());
                    failed_clients.push_back(id);
                }
            }
        }

        for (size_t i = 0; i < failed_clients.size(); ++i) {
            removeClient(failed_clients[i]);
        }
        
        if (success_count > 0) {
            std::ostringstream oss;
            oss << "Broadcasted message to " << success_count << " clients, message size: " << message.length() << " bytes";
            Logger::debug(oss.str());
        }
    }

    size_t getClientCount() const {
        std::lock_guard<std::mutex> lock(clients_mtx_);
        return clients_.size();
    }
};

// === 编码/解码工具函数 ===
constexpr float PI = 3.1415f;

uint16_t encode_gain(int val) {
    if (val < -10000) val = -10000;
    if (val > 10000) val = 10000;
    return (uint16_t)(((int32_t)val + 10000) * 65535 / 20000);
}

int decode_gain(uint16_t code) {
    return (int)(((int32_t)code * 20000) / 65535 - 10000);
}

uint16_t encode_position(float pos) {
    if (pos < -PI) pos = -PI;
    if (pos > PI) pos = PI;
    return (uint16_t)(((pos + PI) * 65535.0f) / (2.0f * PI));
}

float decode_position(uint16_t code) {
    return ((float)code) * 2.0f * PI / 65535.0f - PI;
}

uint16_t encode_current(int val) {
    if (val < -1000) val = -1000;
    if (val > 1000) val = 1000;
    return (uint16_t)(((int32_t)val + 1000) * 65535 / 2000);
}

int decode_current(uint16_t code) {
    return (int)(((int32_t)code * 2000) / 65535 - 1000);
}

// === JSON序列化函数 ===
json joint_status_to_json(const JointParamsUplink& j, int id) {
    return {
        {"id", id},
        {"position", decode_position(j.position)},
        {"Iq", decode_current(j.Iq)},
        {"Id", decode_current(j.Id)},
        {"default1", j.default1},
        {"default2", j.default2},
        {"default3", j.default3},
        {"default4", j.default4},
        {"default5", j.default5},
        {"default6", j.default6},
        {"default7", j.default7}
    };
}

json ydds_to_json(const YddsData& y) {
    return {
        {"nf", y.nf},
        {"tf", y.tf},
        {"tfDir", y.tfDir}
    };
}

json sensor_to_json(const SensorParam& s) {
    json j;
    j["id"] = s.id;
    j["channelData"] = json::array();
    for (int i = 0; i < 11; ++i) j["channelData"].push_back(s.channelData[i]);
    j["yddsData"] = json::array();
    for (int i = 0; i < 2; ++i) j["yddsData"].push_back(ydds_to_json(s.yddsData[i]));
    j["sProxData"] = s.sProxData;
    return j;
}

json palm_sensor_to_json(const SensorParam_Palm_t& s) {
    json j;
    j["id"] = s.sensorID;
    j["channelData"] = json::array();
    for (int i = 0; i < 13; ++i) j["channelData"].push_back(s.channelData[i]);
    j["yddsData"] = json::array();
    for (int i = 0; i < 2; ++i) j["yddsData"].push_back(ydds_to_json(s.yddsData[i]));
    j["sProxData"] = json::array();
    for (int i = 0; i < 2; ++i) j["sProxData"].push_back(s.sProxData[i]);
    return j;
}

json make_status_json(const RS485Uplink& status) {
    json j;
    j["type"] = "status";
    j["timestamp"] = DexterousHand::getCurrentTimeString();
    j["joints"] = json::array();
    for (int i = 0; i < 21; ++i) {
        j["joints"].push_back(joint_status_to_json(status.jointParams[i], i));
    }
    j["sensors"] = json::array();
    for (int i = 0; i < 18; ++i) {
        j["sensors"].push_back(sensor_to_json(status.sensorParams[i]));
    }
    j["palmSensor"] = palm_sensor_to_json(status.SensorParam_Palm);
    return j;
}

// === JSON解析函数 ===
bool parse_set_joints(const json& j, JointCommandsArray& out) {
    if (!j.contains("joints") || !j["joints"].is_array() || j["joints"].size() != 21) {
        Logger::error("Invalid joints array format");
        return false;
    }
    
    try {
        for (int i = 0; i < 21; ++i) {
            const auto& src = j["joints"][i];
            out.joints[i].kp = encode_gain(src.value("kp", 0));
            out.joints[i].ki = encode_gain(src.value("ki", 0));
            out.joints[i].kd = encode_gain(src.value("kd", 0));
            out.joints[i].position = encode_position(src.value("position", 0.0f));
            out.joints[i].current_max = encode_current(src.value("current_max", 0));
            out.joints[i].mode = src.value("mode", 0);
            out.joints[i].res0 = src.value("res0", 0);
            out.joints[i].res1 = src.value("res1", 0);
            out.joints[i].res2 = src.value("res2", 0);
            out.joints[i].res3 = src.value("res3", 0);
        }
        return true;
    } catch (const std::exception& e) {
        std::ostringstream oss;
        oss << "Failed to parse joint parameters: " << e.what();
        Logger::error(oss.str());
        return false;
    }
}

// === 改进的WebSocket服务器类 ===
class ImprovedWebSocketServer {
private:
    std::unique_ptr<DexterousHand> hand_;
    std::string port_;
    
    std::atomic<bool> running_;
    std::thread control_thread_;
    std::thread broadcast_thread_;
    std::thread stats_thread_;
    
    ThreadSafeBuffer<JointCommandsArray> joint_commands_;
    ThreadSafeBuffer<RS485Uplink> latest_status_;
    
    SafeClientManager client_manager_;
    HealthMonitor health_monitor_;
    
    std::map<crow::websocket::connection*, uint64_t> conn_to_id_;
    std::mutex conn_map_mtx_;
    
    std::atomic<bool> use_fake_data_{false};
    std::atomic<int> fake_data_counter_{0};

public:
    ImprovedWebSocketServer(const std::string& port) : port_(port), running_(true) {
        hand_ = std::make_unique<DexterousHand>(port, 2000000);
        hand_->setHandId(0);
        
        JointCommandsArray default_joints;
        for (int i = 0; i < 21; ++i) {
            default_joints.joints[i] = {};
            default_joints.joints[i].kp = encode_gain(100);
            default_joints.joints[i].mode = 3;
        }
        joint_commands_.set(default_joints);
        
        Logger::info("WebSocket server initialized");
    }

    ~ImprovedWebSocketServer() {
        stop();
    }

    bool initialize() {
        if (use_fake_data_) {
            Logger::info("Fake data mode enabled - Skipping hardware initialization");
            return true;
        }
        
        Logger::info("Initializing dexterous hand hardware...");
        if (!hand_->initialize()) {
            Logger::error("Failed to initialize dexterous hand");
            return false;
        }
        
        if (!hand_->startRealTimeControl(50)) {
            Logger::error("Failed to start real-time control");
            return false;
        }
        
        Logger::info("Dexterous hand initialized successfully");
        return true;
    }

    void start() {
        running_ = true;
        
        control_thread_ = std::thread([this]() {
            controlLoop();
        });
        
        broadcast_thread_ = std::thread([this]() {
            broadcastLoop();
        });
        
        stats_thread_ = std::thread([this]() {
            statsLoop();
        });
        
        Logger::info("All worker threads started");
    }

    void stop() {
        running_ = false;
        
        if (control_thread_.joinable()) control_thread_.join();
        if (broadcast_thread_.joinable()) broadcast_thread_.join();
        if (stats_thread_.joinable()) stats_thread_.join();
        
        if (hand_) {
            hand_->close();
        }
        
        Logger::info("All threads stopped");
    }

    void handleClientConnect(crow::websocket::connection& conn) {
        uint64_t client_id = client_manager_.addClient(&conn);
        std::lock_guard<std::mutex> lock(conn_map_mtx_);
        conn_to_id_[&conn] = client_id;
    }

    void handleClientDisconnect(crow::websocket::connection& conn) {
        std::lock_guard<std::mutex> lock(conn_map_mtx_);
        auto it = conn_to_id_.find(&conn);
        if (it != conn_to_id_.end()) {
            client_manager_.removeClient(it->second);
            conn_to_id_.erase(it);
        }
    }

    void handleClientMessage(crow::websocket::connection& conn, const std::string& data) {
        std::lock_guard<std::mutex> lock(conn_map_mtx_);
        auto it = conn_to_id_.find(&conn);
        uint64_t client_id = (it != conn_to_id_.end()) ? it->second : 0;
        
        std::ostringstream oss;
        oss << "Received message from client " << client_id << ", size: " << data.length() << " bytes";
        Logger::debug(oss.str());
        
        try {
            auto j = json::parse(data);
            if (j.value("type", "") == "set_joints") {
                JointCommandsArray new_joints;
                if (parse_set_joints(j, new_joints)) {
                    joint_commands_.set(new_joints);
                    std::ostringstream success_oss;
                    success_oss << "Successfully processed joint control command from client " << client_id;
                    Logger::info(success_oss.str());
                } else {
                    std::ostringstream err_oss;
                    err_oss << "Failed to parse joint command from client " << client_id;
                    Logger::error(err_oss.str());
                }
            } else {
                std::ostringstream warn_oss;
                warn_oss << "Unknown message type '" << j.value("type", "") << "' from client " << client_id;
                Logger::warn(warn_oss.str());
            }
        } catch (const std::exception& e) {
            std::ostringstream err_oss;
            err_oss << "JSON parse error from client " << client_id << ": " << e.what();
            Logger::error(err_oss.str());
        }
    }

    void setFakeDataMode(bool enable) {
        use_fake_data_ = enable;
        if (enable) {
            Logger::info("Fake data mode enabled - Will generate simulated hand data");
        } else {
            Logger::info("Fake data mode disabled - Using real hardware data");
        }
    }
    
    bool isFakeDataMode() const {
        return use_fake_data_;
    }

private:
    void controlLoop() {
        auto period = std::chrono::microseconds(20000);
        auto next_time = std::chrono::steady_clock::now();
        
        Logger::info("Control thread started, frequency: 50Hz");
        if (use_fake_data_) {
            Logger::info("Control thread running in fake data mode");
        }
        
        while (running_.load()) {
            try {
                if (use_fake_data_) {
                    RS485Uplink fake_status = generateFakeHandStatus();
                    latest_status_.set(fake_status);
                    health_monitor_.reportSerialSuccess();
                    
                    auto joints_array = joint_commands_.get();
                } else {
                    auto joints_array = joint_commands_.get();
                    
                    for (int i = 0; i < 21; ++i) {
                        hand_->setJointParams(i, joints_array.joints[i]);
                    }
                    
                    RS485Uplink status = hand_->getHandStatus();
                    latest_status_.set(status);
                    
                    health_monitor_.reportSerialSuccess();
                }
                
            } catch (const std::exception& e) {
                std::ostringstream oss;
                oss << "Control thread error: " << e.what();
                Logger::error(oss.str());
                health_monitor_.reportSerialError();
            }
            
            next_time += period;
            std::this_thread::sleep_until(next_time);
        }
        
        Logger::info("Control thread stopped");
    }

    void broadcastLoop() {
        auto period = std::chrono::milliseconds(20);
        
        Logger::info("Broadcast thread started, frequency: 50Hz");
        
        while (running_.load()) {
            try {
                if (latest_status_.wait_for_data(std::chrono::milliseconds(50))) {
                    auto status = latest_status_.get();
                    auto json_msg = make_status_json(status);
                    auto msg_str = json_msg.dump();
                    
                    client_manager_.broadcast(msg_str);
                    health_monitor_.reportMessageSent();
                }
            } catch (const std::exception& e) {
                std::ostringstream oss;
                oss << "Broadcast thread error: " << e.what();
                Logger::error(oss.str());
            }
            
            std::this_thread::sleep_for(period);
        }
        
        Logger::info("Broadcast thread stopped");
    }

    void statsLoop() {
        while (running_.load()) {
            std::this_thread::sleep_for(std::chrono::seconds(10));
            
            if (running_.load()) {
                health_monitor_.printStats();
                std::ostringstream oss;
                oss << "Active WebSocket clients: " << client_manager_.getClientCount();
                Logger::info(oss.str());
            }
        }
        
        Logger::info("Stats thread stopped");
    }

    RS485Uplink generateFakeHandStatus() {
        RS485Uplink fake_status = {};
        
        auto now = std::chrono::steady_clock::now();
        auto time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();
        int counter = fake_data_counter_++;
        
        for (int i = 0; i < 21; ++i) {
            float phase = (float)i * 0.3f + (float)time_ms * 0.001f;
            float position = sinf(phase) * PI * 0.5f;
            
            int iq_current = (int)(sinf(phase + 1.0f) * 200.0f + (rand() % 20 - 10));
            int id_current = (int)(cosf(phase) * 100.0f + (rand() % 10 - 5));
            
            fake_status.jointParams[i].position = encode_position(position);
            fake_status.jointParams[i].Iq = encode_current(iq_current);
            fake_status.jointParams[i].Id = encode_current(id_current);
            
            fake_status.jointParams[i].default1 = counter % 1000;
            fake_status.jointParams[i].default2 = i * 10;
            fake_status.jointParams[i].default3 = (counter + i) % 256;
            fake_status.jointParams[i].default4 = 0;
            fake_status.jointParams[i].default5 = 0;
            fake_status.jointParams[i].default6 = 0;
            fake_status.jointParams[i].default7 = 0;
        }
        
        for (int i = 0; i < 18; ++i) {
            fake_status.sensorParams[i].id = i;
            
            for (int ch = 0; ch < 11; ++ch) {
                float base_value = 100.0f + sinf((float)(time_ms + i * 100 + ch * 50) * 0.0001f) * 50.0f;
                fake_status.sensorParams[i].channelData[ch] = (uint32_t)(base_value + (rand() % 50));
            }
            
            for (int j = 0; j < 2; ++j) {
                fake_status.sensorParams[i].yddsData[j].nf = (uint32_t)(500 + sinf((float)(time_ms + i * 200 + j * 100) * 0.0002f) * 200);
                fake_status.sensorParams[i].yddsData[j].tf = (uint32_t)(300 + cosf((float)(time_ms + i * 150 + j * 75) * 0.0003f) * 150);
                fake_status.sensorParams[i].yddsData[j].tfDir = (counter + i + j) % 2;
            }
            
            fake_status.sensorParams[i].sProxData = (uint32_t)(200 + sinf((float)(time_ms + i * 300) * 0.0001f) * 100);
        }

        fake_status.SensorParam_Palm.sensorID = 30;
        for (int ch = 0; ch < 13; ++ch) {
            float base_value = 200.0f + sinf((float)(time_ms + ch * 80) * 0.0001f) * 80.0f;
            fake_status.SensorParam_Palm.channelData[ch] = (uint32_t)(base_value + (rand() % 30));
        }
        for (int j = 0; j < 2; ++j) {
            fake_status.SensorParam_Palm.yddsData[j].nf = (uint32_t)(600 + sinf((float)(time_ms + j * 120) * 0.0002f) * 250);
            fake_status.SensorParam_Palm.yddsData[j].tf = (uint32_t)(400 + cosf((float)(time_ms + j * 90) * 0.0003f) * 180);
            fake_status.SensorParam_Palm.yddsData[j].tfDir = (counter + j) % 2;
            fake_status.SensorParam_Palm.sProxData[j] = (uint32_t)(250 + sinf((float)(time_ms + j * 200) * 0.0001f) * 120);
        }
        
        return fake_status;
    }
};

int main(int argc, char* argv[]) {
#ifdef _WIN32
    SetConsoleOutputCP(CP_UTF8);
    SetConsoleCP(CP_UTF8);
#endif

    std::string port;
    bool use_fake = false;

    if (argc > 1) {
        port = argv[1];
        Logger::info("Starting XHand WebSocket Server on serial port " + port + "...");
    } else {
        Logger::info("No serial port specified, using fake data mode.");
        use_fake = true;
        port = "";
    }

    ImprovedWebSocketServer server(port);
    server.setFakeDataMode(use_fake);

    if (!server.initialize()) {
        Logger::error("Server initialization failed");
        return 1;
    }

    server.start();

    crow::SimpleApp app;

    CROW_ROUTE(app, "/ws").websocket()
    .onopen([&server](crow::websocket::connection& conn) {
        server.handleClientConnect(conn);
    })
    .onclose([&server](crow::websocket::connection& conn, const std::string& reason) {
        server.handleClientDisconnect(conn);
    })
    .onmessage([&server](crow::websocket::connection& conn, const std::string& data, bool is_binary) {
        server.handleClientMessage(conn, data);
    });

    Logger::info("WebSocket server starting on port 12345");
    Logger::info("Press Ctrl+C to exit");

    app.port(12345).multithreaded().run();

    server.stop();
    Logger::info("Server shutdown complete");

    return 0;
}